
import torch
import torch.nn as nn
import torchvision

class AlexNet(nn.Module):
  def __init__(self,num_classes=1000):
    super(AlexNet,self).__init__()
    self.feature_extraction = nn.Sequential(
      nn.Conv2d(in_channels=3,out_channels=96,kernel_size=11,stride=4,padding=2,bias=False),
      nn.ReLU(inplace=True),
      nn.MaxPool2d(kernel_size=3,stride=2,padding=0),
      nn.Conv2d(in_channels=96,out_channels=192,kernel_size=5,stride=1,padding=2,bias=False),
      nn.ReLU(inplace=True),
      nn.MaxPool2d(kernel_size=3,stride=2,padding=0),
      nn.Conv2d(in_channels=192,out_channels=384,kernel_size=3,stride=1,padding=1,bias=False),
      nn.ReLU(inplace=True),
      nn.Conv2d(in_channels=384,out_channels=256,kernel_size=3,stride=1,padding=1,bias=False),
      nn.ReLU(inplace=True),
      nn.Conv2d(in_channels=256,out_channels=256,kernel_size=3,stride=1,padding=1,bias=False),
      nn.ReLU(inplace=True),
      nn.MaxPool2d(kernel_size=3, stride=2, padding=0),
    )
    self.classifier = nn.Sequential(
      nn.Dropout(p=0.5),
      nn.Linear(in_features=256*6*6,out_features=4096),
      nn.ReLU(inplace=True),
      nn.Dropout(p=0.5),
      nn.Linear(in_features=4096, out_features=4096),
      nn.ReLU(inplace=True),
      nn.Linear(in_features=4096, out_features=num_classes),
    )
  def forward(self,x):
    x = self.feature_extraction(x)
    # x = x.view(x.size(0),256*6*6)
    # x = self.classifier(x)
    return x

def Alex():
  model = AlexNet()
  # pretrained_dict = torch.load("model_data/alexnet-owt-4df8aa71.pth")

  # model_dict = model.state_dict()
  # # 1. filter out unnecessary keys 筛选出与自己模型Key相同的权重
  # pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
  # # 2. overwrite entries in the existing state dict
  # model_dict.update(pretrained_dict)

  model.load_state_dict(pretrained_dict, strict=False)
  return model


if __name__ == "__main__":
    model = AlexNet()
    # m = torch.load("../model_data/alexnet-owt-4df8aa71.pth")

    # pretrained_dict = torch.load("../model_data/alexnet-owt-4df8aa71.pth")

    # model_dict = model.state_dict()
    # # 1. filter out unnecessary keys 筛选出与自己模型Key相同的权重
    # pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
    # # 2. overwrite entries in the existing state dict
    # model_dict.update(pretrained_dict)

    # model.load_state_dict(pretrained_dict,strict=False)
    from ptflops import get_model_complexity_info
    flops, params = get_model_complexity_info(model, (3, 224, 224), as_strings=True,
                                              print_per_layer_stat=True)  # 不用写batch_size大小，默认batch_size=1
    print('Flops:  ' + flops)
    print('Params: ' + params)

    # print(m)
    # model.load_state_dict(torch.load("../model_data/alexnet-owt-4df8aa71.pth"),strict=False)
    # from torchsummary import summary
    # summary(model, input_size=[(3, 224, 224)], batch_size=2)